package com.heiye.blog.ai.model.openrouter;

import cn.hutool.core.collection.CollUtil;
import com.heiye.blog.framework.common.util.JsonUtils;
import io.micrometer.observation.ObservationRegistry;
import lombok.AllArgsConstructor;
import lombok.Builder;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.MessageType;
import org.springframework.ai.chat.metadata.*;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.model.Generation;
import org.springframework.ai.chat.observation.ChatModelObservationConvention;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.model.ModelOptionsUtils;
import org.springframework.ai.model.tool.ToolCallingManager;
import org.springframework.ai.model.tool.ToolExecutionEligibilityPredicate;
import org.springframework.ai.support.UsageCalculator;
import org.springframework.retry.support.RetryTemplate;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.util.Collection;
import java.util.List;
import java.util.Objects;

/**
 * @author: heiye
 * @date: 2025/11/11 下午10:01
 * @version: v1.0.0
 * @description: OpenRouterChatModel
 */
@Data
@Slf4j
@Builder
@AllArgsConstructor
public class OpenRouterChatModel implements ChatModel {

    private final OpenRouterApi openRouterApi;

    private final OpenRouterChatOptions defaultOptions;

    private final ObservationRegistry observationRegistry;

    public final RetryTemplate retryTemplate;

    private final ToolCallingManager toolCallingManager;

    private final ToolExecutionEligibilityPredicate toolExecutionEligibilityPredicate;

    private ChatModelObservationConvention observationConvention;

    /**
     * 在 Spring AI 中，这个方法是把 Prompt 转为对应调用 Api 的 Prompt, 并把工具参数赋好值
     * 如由 ChatOptions 接口类型转化为 DeepSeek 的 ChatOptions (DeepSeekChatOptions) 类
     *
     * @param prompt
     * @return
     */
    public Prompt buildRequestPrompt(Prompt prompt) {
        OpenRouterChatOptions runtimeOptions = null;

        // 先创建对应 Api 的 ChatOptions 类似于 DeepSeekChatOptions，这里用 OpenRouterChatOptions 代替
        // 判断 Prompt 里的 ChatOptions 是否为空
        // 不为空，判断里面的 ChatOptions 类型，如果是 ToolCallingChatOptions 类型，
        //      也就是调用工具，把里面的 ToolCallingChatOptions 工具类参数则 bean 拷贝为 OpenRouterChatOptions 工参数，
        //      否则拷贝为普通的聊天参数 (OpenRouterChatOptions)
        if (prompt.getOptions() != null) {
            // 这里不用判断工具调用，因为 OpenRouter 暂时不支持工具调用
            runtimeOptions = (OpenRouterChatOptions) ModelOptionsUtils.copyToTarget(prompt.getOptions(), ChatOptions.class, OpenRouterChatOptions.class);
        }

        // 这里 bean 拷贝直接把默认值赋值 ChatOptions 里，转为 OpenRouterChatOptions 类型获得 requestOptions 参数
        OpenRouterChatOptions requestOptions = (OpenRouterChatOptions) ModelOptionsUtils.merge(runtimeOptions, this.defaultOptions, OpenRouterChatOptions.class);

        // 之后就是判断 requestOptions 是否空来塞调用工具的参数
        // 空就给你塞默认值
        // 没空就给你塞原本 ChatOptions 的工具参数
//        if (runtimeOptions != null) {
//            requestOptions.setInternalToolExecutionEnabled((Boolean) ModelOptionsUtils.mergeOption(runtimeOptions.getInternalToolExecutionEnabled(), this.defaultOptions.getInternalToolExecutionEnabled()));
//            requestOptions.setToolNames(ToolCallingChatOptions.mergeToolNames(runtimeOptions.getToolNames(), this.defaultOptions.getToolNames()));
//            requestOptions.setToolCallbacks(ToolCallingChatOptions.mergeToolCallbacks(runtimeOptions.getToolCallbacks(), this.defaultOptions.getToolCallbacks()));
//            requestOptions.setToolContext(ToolCallingChatOptions.mergeToolContext(runtimeOptions.getToolContext(), this.defaultOptions.getToolContext()));
//        } else {
//            requestOptions.setInternalToolExecutionEnabled(this.defaultOptions.getInternalToolExecutionEnabled());
//            requestOptions.setToolNames(this.defaultOptions.getToolNames());
//            requestOptions.setToolCallbacks(this.defaultOptions.getToolCallbacks());
//            requestOptions.setToolContext(this.defaultOptions.getToolContext());
//        }
//
//        ToolCallingChatOptions.validateToolCallbacks(requestOptions.getToolCallbacks());

        // 这里不处理，因为 OpenRouter 暂时不支持工具调用
        return new Prompt(prompt.getInstructions(), requestOptions);
    }

    /**
     * 流式处理
     *
     * @param prompt
     * @return
     */
    public Flux<ChatResponse> internalStream(Prompt prompt, ChatResponse previousChatResponse) {
        return Flux.deferContextual((contextView) -> {
            // 创建请求体
            OpenRouterApi.ChatCompletionRequest request = this.createRequest(prompt, true);

            // 发送请求获取到流
            Flux<OpenRouterApi.ChatCompletionChunk> completionChunks = this.openRouterApi.chatCompletionStream(request);

            // 处理值转化为 OpenRouterChatResponse 类型
            Flux<ChatResponse> chatResponseFlux = completionChunks
                    // 顺序处理
                    .switchMap(chunk -> Mono.just(chunk).map(chunk2 -> {
                        // 获取到 OpenRouter 生成数据
                        List<Generation> generations = chunk2.choices().stream()
                                .map(this::buildGeneration).toList();

                        // 获取到该块的 token 使用数据
                        OpenRouterApi.Usage usage = chunk2.usage();
                        Usage currentUsage = usage != null ? this.getDefaultUsage(usage) : new EmptyUsage();

                        Usage cumulativeUsage = UsageCalculator.getCumulativeUsage(currentUsage, previousChatResponse);
                        // 进行反参
                        return new ChatResponse(generations, this.from(chunk2, cumulativeUsage));
                    }));

            // Spring AI 在最后会进行一次聚合，对 token 用量、文本元数据等聚合
            // 在这里 OpenRouter 会在最后得 chunk 中会返回最后的用量数据
            // 文本元数据什么的都给一个上下文对象拿走了，还没输出一点都没处理
            // 文本元数据，我会在 advisor 聚合写入数据库
            return chatResponseFlux;

        });
    }

    /**
     * 构建 OpenRouter 生成
     *
     * @param choice
     * @return
     */
    public Generation buildGeneration(OpenRouterApi.ChatCompletionChunk.Choice choice) {
        // 创建生成元数据
        String finishReason = choice.finishReason() != null ?
                choice.finishReason().name() :
                // 当 finishReason 不存在时，使用 nativeFinishReason，如果没有就是没有了
                choice.nativeFinishReason() != null ? choice.nativeFinishReason().name()
                        : null;

        // 创建 Builder
        ChatGenerationMetadata.Builder generationMetadataBuilder = ChatGenerationMetadata.builder().finishReason(finishReason);

        // 获取推理内容
        String reasoningContent = choice.delta().reasoning();

        // 获取 choice 的文本内容
        String content = choice.delta().content();

        // 构建 AI 助手消息, 不使用 builder 是因为有其他参数，剩余的参数由父类默认处理就好
        OpenRouterAssistantMessage openRouterAssistantMessage = new OpenRouterAssistantMessage(content, reasoningContent);

        // 反参
        return new Generation(openRouterAssistantMessage, generationMetadataBuilder.build());
    }

    /**
     * 创建  OpenRouter 请求体
     *
     * @param prompt
     * @param stream
     * @return
     */
    OpenRouterApi.ChatCompletionRequest createRequest(Prompt prompt, Boolean stream) {

        // 获取聊天消息
        List<Message> messages = prompt.getInstructions();
        List<OpenRouterApi.ChatCompletionRequestMessage> chatCompletionRequestMessages = messages.stream().map((message) -> {

            // 在 spring ai 中 他会判断 消息类型，如果不是用户消息和系统消息
            // 就判断是否是助手消息或工具消息
            if (message.getMessageType() == MessageType.ASSISTANT) {
                // 如果是助手消息: 转化为 助手消息 (AssistantMessage)
                //               创建工具调用变量列表
                //               如果这个助手消息调用了工具，获取该列表
                //               将这个工具调用变量添加到工具调用列表中
                // 这里暂时不支持工具调用
                //                 如果这个助手消息没调用工具，那么就返回 AssistantMessage
                // 这里调用 OpenRouterChatCompletionMessage 里的 OpenRouterChatCompletionTextMessage 来构建用户文本消息
                return List.of(OpenRouterApi.ChatCompletionRequestMessage.OpenRouterChatCompletionTextMessage(
                        // 消息角色
                        OpenRouterApi.Role.valueOf(message.getMessageType().name()),
                        // 文本内容
                        message.getText()));
            }

            // 如果是工具消息: 转化为工具消息 (ToolResponseMessage)
            //                遍历 toolMessage 进行断言判断，判断你这个 Response 是否为空
            //                对 toolMessage 遍历构建传输的工具消息

            return List.of(OpenRouterApi.ChatCompletionRequestMessage.OpenRouterChatCompletionTextMessage(OpenRouterApi.Role.valueOf(message.getMessageType().name()), message.getText()));

        }).flatMap(Collection::stream).toList();

        // 创建请求
        OpenRouterApi.ChatCompletionRequest request = new OpenRouterApi.ChatCompletionRequest(chatCompletionRequestMessages, stream);
        // 获取 OpenRouter 聊天选项
        OpenRouterChatOptions requestOptions = (OpenRouterChatOptions) prompt.getOptions();
        // 进行参数转移
        request = OpenRouterRequestAdapter.adaptOptionsToRequest(requestOptions, request);

        return request;
    }

    /**
     * 获取传输的Usage
     *
     * @param usage
     * @return
     */
    private DefaultUsage getDefaultUsage(OpenRouterApi.Usage usage) {
        return new DefaultUsage(usage.promptTokens(), usage.completionTokens(), usage.totalTokens(), usage);
    }

    /**
     * 构建 ChatResponseMetadata
     *
     * @param result
     * @param usage
     * @return
     */
    private ChatResponseMetadata from(OpenRouterApi.ChatCompletionChunk result, Usage usage) {
        ChatResponseMetadata.Builder builder = ChatResponseMetadata.builder()
                // 构建 id
                .id(result.id() != null ? result.id() : "")
                // 构建 usage
                .usage(usage)
                // 构建 model
                .model(result.model() != null ? result.model() : "")
                // 构建 metadata
                .keyValue("created", result.created() != null ? result.created() : 0L)
                .keyValue("system-fingerprint", result.systemFingerprint() != null ? result.systemFingerprint() : "");

        return builder.build();
    }

    /**
     * 设置 ObservationConvention
     *
     * @param observationConvention
     */
    public void setObservationConvention(ChatModelObservationConvention observationConvention) {
        this.observationConvention = observationConvention;
    }
    @Override
    public String call(String message) {
        return ChatModel.super.call(message);
    }

    @Override
    public String call(Message... messages) {
        return ChatModel.super.call(messages);
    }

    @Override
    public ChatResponse call(Prompt prompt) {
        return null;
    }

    @Override
    public ChatOptions getDefaultOptions() {
        return ChatModel.super.getDefaultOptions();
    }

    @Override
    public Flux<String> stream(String message) {
        return ChatModel.super.stream(message);
    }

    @Override
    public Flux<String> stream(Message... messages) {
        return ChatModel.super.stream(messages);
    }

    @Override
    public Flux<ChatResponse> stream(Prompt prompt) {
        // 构建请求体提示词
        Prompt requestPrompt = this.buildRequestPrompt(prompt);
        // 调用 internalStream 方法，获取流式数据
        return this.internalStream(requestPrompt, null);
    }
}
